Skip to content

[PGO] Add ProfileInjector and ProfileVerifier passes #147388

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

mtrofin
Copy link
Member

@mtrofin mtrofin commented Jul 7, 2025

Adding 2 passes, one to inject MD_prof and one to check its presence. A subsequent patch will add these (similar to debugify) to opt (and, eventually, a variant of this, to llc)

Tracking issue: #147390

Copy link
Member Author

mtrofin commented Jul 7, 2025

This stack of pull requests is managed by Graphite. Learn more about stacking.

@mtrofin mtrofin changed the title pass validator [PGO] Add ProfileInjector and ProfileVerifier passes Jul 7, 2025
@mtrofin mtrofin marked this pull request as ready for review July 7, 2025 20:30
@llvmbot llvmbot added PGO Profile Guided Optimizations llvm:transforms labels Jul 7, 2025
@llvmbot
Copy link
Member

llvmbot commented Jul 7, 2025

@llvm/pr-subscribers-pgo

Author: Mircea Trofin (mtrofin)

Changes

Adding 2 passes, one to inject MD_prof and one to check its presence. A subsequent patch will add these (similar to debugify) to opt (and, eventually, a variant of this, to llc)

Tracking issue: #147390


Full diff: https://github.com/llvm/llvm-project/pull/147388.diff

8 Files Affected:

  • (added) llvm/include/llvm/Transforms/Utils/ProfileVerify.h (+36)
  • (modified) llvm/lib/Passes/PassBuilder.cpp (+1)
  • (modified) llvm/lib/Passes/PassRegistry.def (+2)
  • (modified) llvm/lib/Transforms/Utils/CMakeLists.txt (+1)
  • (added) llvm/lib/Transforms/Utils/ProfileVerify.cpp (+113)
  • (added) llvm/test/Transforms/PGOProfile/prof-verify-as-needed.ll (+19)
  • (added) llvm/test/Transforms/PGOProfile/prof-verify-existing.ll (+21)
  • (added) llvm/test/Transforms/PGOProfile/prof-verify.ll (+19)
diff --git a/llvm/include/llvm/Transforms/Utils/ProfileVerify.h b/llvm/include/llvm/Transforms/Utils/ProfileVerify.h
new file mode 100644
index 0000000000000..88942a73474d4
--- /dev/null
+++ b/llvm/include/llvm/Transforms/Utils/ProfileVerify.h
@@ -0,0 +1,36 @@
+//===- ProfileVerify.h - Verify profile info for testing ----0-----*-C++-*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Inject profile information, as part of tests, to verify passes don't
+// accidentally drop it.
+//
+//===----------------------------------------------------------------------===//
+#ifndef LLVM_TRANSFORMS_UTILS_PROFILEVERIFY_H
+#define LLVM_TRANSFORMS_UTILS_PROFILEVERIFY_H
+
+#include "llvm/IR/Analysis.h"
+#include "llvm/IR/PassManager.h"
+
+namespace llvm {
+/// Inject MD_prof metadata where it's missing. Used for testing that passes
+/// don't accidentally drop this metadata.
+class ProfileInjectorPass : public PassInfoMixin<ProfileInjectorPass> {
+public:
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &FAM);
+};
+
+/// Checks that MD_prof is present on every instruction that supports it. Used
+/// in conjunction with the ProfileInjectorPass. MD_prof "unknown" is considered
+/// valid (i.e. !{!"unknown"})
+class ProfileVerifierPass : public PassInfoMixin<ProfileVerifierPass> {
+public:
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &FAM);
+};
+
+} // namespace llvm
+#endif
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
index 874fce05841e2..70c0e999cbb95 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -357,6 +357,7 @@
 #include "llvm/Transforms/Utils/MoveAutoInit.h"
 #include "llvm/Transforms/Utils/NameAnonGlobals.h"
 #include "llvm/Transforms/Utils/PredicateInfo.h"
+#include "llvm/Transforms/Utils/ProfileVerify.h"
 #include "llvm/Transforms/Utils/RelLookupTableConverter.h"
 #include "llvm/Transforms/Utils/StripGCRelocates.h"
 #include "llvm/Transforms/Utils/StripNonLineTableDebugInfo.h"
diff --git a/llvm/lib/Passes/PassRegistry.def b/llvm/lib/Passes/PassRegistry.def
index dd3dab3425975..dfe233f15a149 100644
--- a/llvm/lib/Passes/PassRegistry.def
+++ b/llvm/lib/Passes/PassRegistry.def
@@ -517,6 +517,8 @@ FUNCTION_PASS("print<regions>", RegionInfoPrinterPass(errs()))
 FUNCTION_PASS("print<scalar-evolution>", ScalarEvolutionPrinterPass(errs()))
 FUNCTION_PASS("print<stack-safety-local>", StackSafetyPrinterPass(errs()))
 FUNCTION_PASS("print<uniformity>", UniformityInfoPrinterPass(errs()))
+FUNCTION_PASS("prof-inject", ProfileInjectorPass())
+FUNCTION_PASS("prof-verify", ProfileVerifierPass())
 FUNCTION_PASS("reassociate", ReassociatePass())
 FUNCTION_PASS("redundant-dbg-inst-elim", RedundantDbgInstEliminationPass())
 FUNCTION_PASS("reg2mem", RegToMemPass())
diff --git a/llvm/lib/Transforms/Utils/CMakeLists.txt b/llvm/lib/Transforms/Utils/CMakeLists.txt
index 78cad0d253be8..c0bd6d647aad4 100644
--- a/llvm/lib/Transforms/Utils/CMakeLists.txt
+++ b/llvm/lib/Transforms/Utils/CMakeLists.txt
@@ -67,6 +67,7 @@ add_llvm_component_library(LLVMTransformUtils
   MoveAutoInit.cpp
   NameAnonGlobals.cpp
   PredicateInfo.cpp
+  ProfileVerify.cpp
   PromoteMemoryToRegister.cpp
   RelLookupTableConverter.cpp
   ScalarEvolutionExpander.cpp
diff --git a/llvm/lib/Transforms/Utils/ProfileVerify.cpp b/llvm/lib/Transforms/Utils/ProfileVerify.cpp
new file mode 100644
index 0000000000000..e0d220e998548
--- /dev/null
+++ b/llvm/lib/Transforms/Utils/ProfileVerify.cpp
@@ -0,0 +1,113 @@
+//===- ProfileVerify.cpp - Verify profile info for testing ----------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Utils/ProfileVerify.h"
+#include "llvm/ADT/DynamicAPInt.h"
+#include "llvm/ADT/PostOrderIterator.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Analysis/BranchProbabilityInfo.h"
+#include "llvm/Analysis/LoopInfo.h"
+#include "llvm/IR/Analysis.h"
+#include "llvm/IR/Dominators.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/MDBuilder.h"
+#include "llvm/IR/ProfDataUtils.h"
+#include "llvm/Support/BranchProbability.h"
+
+using namespace llvm;
+namespace {
+class ProfileInjector {
+  Function &F;
+  FunctionAnalysisManager &FAM;
+
+public:
+  static bool supportsBranchWeights(const Instruction &I) {
+    return isa<BranchInst>(&I) ||
+
+           isa<SwitchInst>(&I) ||
+
+           isa<IndirectBrInst>(&I) || isa<SelectInst>(&I) ||
+           isa<CallBrInst>(&I);
+  }
+
+  ProfileInjector(Function &F, FunctionAnalysisManager &FAM) : F(F), FAM(FAM) {}
+  bool inject();
+};
+} // namespace
+
+bool ProfileInjector::inject() {
+  auto &BPI = FAM.getResult<BranchProbabilityAnalysis>(F);
+
+  for (auto &BB : F) {
+    if (succ_size(&BB) <= 1)
+      continue;
+    auto *Term = BB.getTerminator();
+    assert(Term);
+    if (Term->getMetadata(LLVMContext::MD_prof) ||
+        !supportsBranchWeights(*Term))
+      continue;
+    SmallVector<BranchProbability> Probs;
+    Probs.reserve(Term->getNumSuccessors());
+    for (auto I = 0U, E = Term->getNumSuccessors(); I < E; ++I)
+      Probs.emplace_back(BPI.getEdgeProbability(&BB, Term->getSuccessor(I)));
+
+    const auto *FirstZeroDenominator =
+        find_if(Probs, [](const BranchProbability &P) {
+          return P.getDenominator() == 0;
+        });
+    assert(FirstZeroDenominator == Probs.end());
+    const auto *FirstNonzeroNumerator =
+        find_if(Probs, [](const BranchProbability &P) {
+          return P.getNumerator() != 0;
+        });
+    assert(FirstNonzeroNumerator != Probs.end());
+    DynamicAPInt LCM(Probs[0].getDenominator());
+    DynamicAPInt GCD(FirstNonzeroNumerator->getNumerator());
+    for (const auto &Prob : drop_begin(Probs)) {
+      if (!Prob.getNumerator())
+        continue;
+      LCM = llvm::lcm(LCM, DynamicAPInt(Prob.getDenominator()));
+      GCD = llvm::lcm(GCD, DynamicAPInt(Prob.getNumerator()));
+    }
+    SmallVector<uint32_t> Weights;
+    Weights.reserve(Term->getNumSuccessors());
+    for (const auto &Prob : Probs) {
+      auto W = Prob.getNumerator() * LCM / GCD;
+      Weights.emplace_back(static_cast<int32_t>((int64_t)W));
+    }
+    setBranchWeights(*Term, Weights, false);
+  }
+  return true;
+}
+
+PreservedAnalyses ProfileInjectorPass::run(Function &F,
+                                           FunctionAnalysisManager &FAM) {
+  ProfileInjector PI(F, FAM);
+  if (!PI.inject())
+    return PreservedAnalyses::all();
+
+  return PreservedAnalyses::none();
+}
+
+PreservedAnalyses ProfileVerifierPass::run(Function &F,
+                                           FunctionAnalysisManager &FAM) {
+  bool Changed = false;
+  for (auto &BB : F)
+    if (succ_size(&BB) >= 2)
+      if (auto *Term = BB.getTerminator())
+        if (ProfileInjector::supportsBranchWeights(*Term)) {
+          if (!Term->getMetadata(LLVMContext::MD_prof)) {
+            F.getContext().emitError("Profile verification failed");
+          } else {
+            Changed = true;
+          }
+        }
+
+  return Changed ? PreservedAnalyses::none() : PreservedAnalyses::all();
+}
diff --git a/llvm/test/Transforms/PGOProfile/prof-verify-as-needed.ll b/llvm/test/Transforms/PGOProfile/prof-verify-as-needed.ll
new file mode 100644
index 0000000000000..19d60d03bd873
--- /dev/null
+++ b/llvm/test/Transforms/PGOProfile/prof-verify-as-needed.ll
@@ -0,0 +1,19 @@
+; Test that prof-inject only injects missing metadata
+
+; RUN: opt -passes=prof-inject %s -S -o - | FileCheck %s
+
+define void @foo(i32 %i) {
+  %c = icmp eq i32 %i, 0
+  br i1 %c, label %yes, label %no, !prof !0
+yes:
+  br i1 %c, label %yes2, label %no
+yes2:
+  ret void
+no:
+  ret void
+}
+
+!0 = !{!"branch_weights", i32 1, i32 2}
+; CHECK: br i1 %c, label %yes, label %no, !prof !0
+; CHECK: !0 = !{!"branch_weights", i32 1, i32 2}
+; CHECK: !1 = !{!"branch_weights", i32 429496729, i32 715827882}
diff --git a/llvm/test/Transforms/PGOProfile/prof-verify-existing.ll b/llvm/test/Transforms/PGOProfile/prof-verify-existing.ll
new file mode 100644
index 0000000000000..d9cc038bc3204
--- /dev/null
+++ b/llvm/test/Transforms/PGOProfile/prof-verify-existing.ll
@@ -0,0 +1,21 @@
+; Test that prof-verify does not modify existing metadata (incl. "unknown")
+
+; RUN: opt -passes=prof-inject %s -S -o - | FileCheck %s
+; RUN: opt -passes=prof-verify %s -S --disable-output
+
+define void @foo(i32 %i) {
+  %c = icmp eq i32 %i, 0
+  br i1 %c, label %yes, label %no, !prof !0
+yes:
+  br i1 %c, label %yes2, label %no, !prof !1
+yes2:
+  ret void
+no:
+  ret void
+}
+
+!0 = !{!"branch_weights", i32 1, i32 2}
+!1 = !{!"unknown"}
+; CHECK: br i1 %c, label %yes, label %no, !prof !0
+; CHECK: !0 = !{!"branch_weights", i32 1, i32 2}
+; CHECK: !1 = !{!"unknown"}
diff --git a/llvm/test/Transforms/PGOProfile/prof-verify.ll b/llvm/test/Transforms/PGOProfile/prof-verify.ll
new file mode 100644
index 0000000000000..f4d2b2e51dc59
--- /dev/null
+++ b/llvm/test/Transforms/PGOProfile/prof-verify.ll
@@ -0,0 +1,19 @@
+; Test prof-inject and prof-verify
+
+; RUN: opt -passes=prof-inject %s -S -o - | FileCheck %s --check-prefix=INJECT
+; RUN: not opt -passes=prof-verify %s -S -o - 2>&1 | FileCheck %s --check-prefix=VERIFY
+; RUN: opt -passes=prof-inject,prof-verify %s --disable-output
+
+define void @foo(i32 %i) {
+  %c = icmp eq i32 %i, 0
+  br i1 %c, label %yes, label %no
+yes:
+  ret void
+no:
+  ret void
+}
+
+; INJECT: br i1 %c, label %yes, label %no, !prof !0
+; INJECT: !0 = !{!"branch_weights", i32 429496729, i32 715827882}
+
+; VERIFY: Profile verification failed
\ No newline at end of file

@llvmbot
Copy link
Member

llvmbot commented Jul 7, 2025

@llvm/pr-subscribers-llvm-transforms

Author: Mircea Trofin (mtrofin)

Changes

Adding 2 passes, one to inject MD_prof and one to check its presence. A subsequent patch will add these (similar to debugify) to opt (and, eventually, a variant of this, to llc)

Tracking issue: #147390


Full diff: https://github.com/llvm/llvm-project/pull/147388.diff

8 Files Affected:

  • (added) llvm/include/llvm/Transforms/Utils/ProfileVerify.h (+36)
  • (modified) llvm/lib/Passes/PassBuilder.cpp (+1)
  • (modified) llvm/lib/Passes/PassRegistry.def (+2)
  • (modified) llvm/lib/Transforms/Utils/CMakeLists.txt (+1)
  • (added) llvm/lib/Transforms/Utils/ProfileVerify.cpp (+113)
  • (added) llvm/test/Transforms/PGOProfile/prof-verify-as-needed.ll (+19)
  • (added) llvm/test/Transforms/PGOProfile/prof-verify-existing.ll (+21)
  • (added) llvm/test/Transforms/PGOProfile/prof-verify.ll (+19)
diff --git a/llvm/include/llvm/Transforms/Utils/ProfileVerify.h b/llvm/include/llvm/Transforms/Utils/ProfileVerify.h
new file mode 100644
index 0000000000000..88942a73474d4
--- /dev/null
+++ b/llvm/include/llvm/Transforms/Utils/ProfileVerify.h
@@ -0,0 +1,36 @@
+//===- ProfileVerify.h - Verify profile info for testing ----0-----*-C++-*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Inject profile information, as part of tests, to verify passes don't
+// accidentally drop it.
+//
+//===----------------------------------------------------------------------===//
+#ifndef LLVM_TRANSFORMS_UTILS_PROFILEVERIFY_H
+#define LLVM_TRANSFORMS_UTILS_PROFILEVERIFY_H
+
+#include "llvm/IR/Analysis.h"
+#include "llvm/IR/PassManager.h"
+
+namespace llvm {
+/// Inject MD_prof metadata where it's missing. Used for testing that passes
+/// don't accidentally drop this metadata.
+class ProfileInjectorPass : public PassInfoMixin<ProfileInjectorPass> {
+public:
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &FAM);
+};
+
+/// Checks that MD_prof is present on every instruction that supports it. Used
+/// in conjunction with the ProfileInjectorPass. MD_prof "unknown" is considered
+/// valid (i.e. !{!"unknown"})
+class ProfileVerifierPass : public PassInfoMixin<ProfileVerifierPass> {
+public:
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &FAM);
+};
+
+} // namespace llvm
+#endif
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
index 874fce05841e2..70c0e999cbb95 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -357,6 +357,7 @@
 #include "llvm/Transforms/Utils/MoveAutoInit.h"
 #include "llvm/Transforms/Utils/NameAnonGlobals.h"
 #include "llvm/Transforms/Utils/PredicateInfo.h"
+#include "llvm/Transforms/Utils/ProfileVerify.h"
 #include "llvm/Transforms/Utils/RelLookupTableConverter.h"
 #include "llvm/Transforms/Utils/StripGCRelocates.h"
 #include "llvm/Transforms/Utils/StripNonLineTableDebugInfo.h"
diff --git a/llvm/lib/Passes/PassRegistry.def b/llvm/lib/Passes/PassRegistry.def
index dd3dab3425975..dfe233f15a149 100644
--- a/llvm/lib/Passes/PassRegistry.def
+++ b/llvm/lib/Passes/PassRegistry.def
@@ -517,6 +517,8 @@ FUNCTION_PASS("print<regions>", RegionInfoPrinterPass(errs()))
 FUNCTION_PASS("print<scalar-evolution>", ScalarEvolutionPrinterPass(errs()))
 FUNCTION_PASS("print<stack-safety-local>", StackSafetyPrinterPass(errs()))
 FUNCTION_PASS("print<uniformity>", UniformityInfoPrinterPass(errs()))
+FUNCTION_PASS("prof-inject", ProfileInjectorPass())
+FUNCTION_PASS("prof-verify", ProfileVerifierPass())
 FUNCTION_PASS("reassociate", ReassociatePass())
 FUNCTION_PASS("redundant-dbg-inst-elim", RedundantDbgInstEliminationPass())
 FUNCTION_PASS("reg2mem", RegToMemPass())
diff --git a/llvm/lib/Transforms/Utils/CMakeLists.txt b/llvm/lib/Transforms/Utils/CMakeLists.txt
index 78cad0d253be8..c0bd6d647aad4 100644
--- a/llvm/lib/Transforms/Utils/CMakeLists.txt
+++ b/llvm/lib/Transforms/Utils/CMakeLists.txt
@@ -67,6 +67,7 @@ add_llvm_component_library(LLVMTransformUtils
   MoveAutoInit.cpp
   NameAnonGlobals.cpp
   PredicateInfo.cpp
+  ProfileVerify.cpp
   PromoteMemoryToRegister.cpp
   RelLookupTableConverter.cpp
   ScalarEvolutionExpander.cpp
diff --git a/llvm/lib/Transforms/Utils/ProfileVerify.cpp b/llvm/lib/Transforms/Utils/ProfileVerify.cpp
new file mode 100644
index 0000000000000..e0d220e998548
--- /dev/null
+++ b/llvm/lib/Transforms/Utils/ProfileVerify.cpp
@@ -0,0 +1,113 @@
+//===- ProfileVerify.cpp - Verify profile info for testing ----------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Utils/ProfileVerify.h"
+#include "llvm/ADT/DynamicAPInt.h"
+#include "llvm/ADT/PostOrderIterator.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Analysis/BranchProbabilityInfo.h"
+#include "llvm/Analysis/LoopInfo.h"
+#include "llvm/IR/Analysis.h"
+#include "llvm/IR/Dominators.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/MDBuilder.h"
+#include "llvm/IR/ProfDataUtils.h"
+#include "llvm/Support/BranchProbability.h"
+
+using namespace llvm;
+namespace {
+class ProfileInjector {
+  Function &F;
+  FunctionAnalysisManager &FAM;
+
+public:
+  static bool supportsBranchWeights(const Instruction &I) {
+    return isa<BranchInst>(&I) ||
+
+           isa<SwitchInst>(&I) ||
+
+           isa<IndirectBrInst>(&I) || isa<SelectInst>(&I) ||
+           isa<CallBrInst>(&I);
+  }
+
+  ProfileInjector(Function &F, FunctionAnalysisManager &FAM) : F(F), FAM(FAM) {}
+  bool inject();
+};
+} // namespace
+
+bool ProfileInjector::inject() {
+  auto &BPI = FAM.getResult<BranchProbabilityAnalysis>(F);
+
+  for (auto &BB : F) {
+    if (succ_size(&BB) <= 1)
+      continue;
+    auto *Term = BB.getTerminator();
+    assert(Term);
+    if (Term->getMetadata(LLVMContext::MD_prof) ||
+        !supportsBranchWeights(*Term))
+      continue;
+    SmallVector<BranchProbability> Probs;
+    Probs.reserve(Term->getNumSuccessors());
+    for (auto I = 0U, E = Term->getNumSuccessors(); I < E; ++I)
+      Probs.emplace_back(BPI.getEdgeProbability(&BB, Term->getSuccessor(I)));
+
+    const auto *FirstZeroDenominator =
+        find_if(Probs, [](const BranchProbability &P) {
+          return P.getDenominator() == 0;
+        });
+    assert(FirstZeroDenominator == Probs.end());
+    const auto *FirstNonzeroNumerator =
+        find_if(Probs, [](const BranchProbability &P) {
+          return P.getNumerator() != 0;
+        });
+    assert(FirstNonzeroNumerator != Probs.end());
+    DynamicAPInt LCM(Probs[0].getDenominator());
+    DynamicAPInt GCD(FirstNonzeroNumerator->getNumerator());
+    for (const auto &Prob : drop_begin(Probs)) {
+      if (!Prob.getNumerator())
+        continue;
+      LCM = llvm::lcm(LCM, DynamicAPInt(Prob.getDenominator()));
+      GCD = llvm::lcm(GCD, DynamicAPInt(Prob.getNumerator()));
+    }
+    SmallVector<uint32_t> Weights;
+    Weights.reserve(Term->getNumSuccessors());
+    for (const auto &Prob : Probs) {
+      auto W = Prob.getNumerator() * LCM / GCD;
+      Weights.emplace_back(static_cast<int32_t>((int64_t)W));
+    }
+    setBranchWeights(*Term, Weights, false);
+  }
+  return true;
+}
+
+PreservedAnalyses ProfileInjectorPass::run(Function &F,
+                                           FunctionAnalysisManager &FAM) {
+  ProfileInjector PI(F, FAM);
+  if (!PI.inject())
+    return PreservedAnalyses::all();
+
+  return PreservedAnalyses::none();
+}
+
+PreservedAnalyses ProfileVerifierPass::run(Function &F,
+                                           FunctionAnalysisManager &FAM) {
+  bool Changed = false;
+  for (auto &BB : F)
+    if (succ_size(&BB) >= 2)
+      if (auto *Term = BB.getTerminator())
+        if (ProfileInjector::supportsBranchWeights(*Term)) {
+          if (!Term->getMetadata(LLVMContext::MD_prof)) {
+            F.getContext().emitError("Profile verification failed");
+          } else {
+            Changed = true;
+          }
+        }
+
+  return Changed ? PreservedAnalyses::none() : PreservedAnalyses::all();
+}
diff --git a/llvm/test/Transforms/PGOProfile/prof-verify-as-needed.ll b/llvm/test/Transforms/PGOProfile/prof-verify-as-needed.ll
new file mode 100644
index 0000000000000..19d60d03bd873
--- /dev/null
+++ b/llvm/test/Transforms/PGOProfile/prof-verify-as-needed.ll
@@ -0,0 +1,19 @@
+; Test that prof-inject only injects missing metadata
+
+; RUN: opt -passes=prof-inject %s -S -o - | FileCheck %s
+
+define void @foo(i32 %i) {
+  %c = icmp eq i32 %i, 0
+  br i1 %c, label %yes, label %no, !prof !0
+yes:
+  br i1 %c, label %yes2, label %no
+yes2:
+  ret void
+no:
+  ret void
+}
+
+!0 = !{!"branch_weights", i32 1, i32 2}
+; CHECK: br i1 %c, label %yes, label %no, !prof !0
+; CHECK: !0 = !{!"branch_weights", i32 1, i32 2}
+; CHECK: !1 = !{!"branch_weights", i32 429496729, i32 715827882}
diff --git a/llvm/test/Transforms/PGOProfile/prof-verify-existing.ll b/llvm/test/Transforms/PGOProfile/prof-verify-existing.ll
new file mode 100644
index 0000000000000..d9cc038bc3204
--- /dev/null
+++ b/llvm/test/Transforms/PGOProfile/prof-verify-existing.ll
@@ -0,0 +1,21 @@
+; Test that prof-verify does not modify existing metadata (incl. "unknown")
+
+; RUN: opt -passes=prof-inject %s -S -o - | FileCheck %s
+; RUN: opt -passes=prof-verify %s -S --disable-output
+
+define void @foo(i32 %i) {
+  %c = icmp eq i32 %i, 0
+  br i1 %c, label %yes, label %no, !prof !0
+yes:
+  br i1 %c, label %yes2, label %no, !prof !1
+yes2:
+  ret void
+no:
+  ret void
+}
+
+!0 = !{!"branch_weights", i32 1, i32 2}
+!1 = !{!"unknown"}
+; CHECK: br i1 %c, label %yes, label %no, !prof !0
+; CHECK: !0 = !{!"branch_weights", i32 1, i32 2}
+; CHECK: !1 = !{!"unknown"}
diff --git a/llvm/test/Transforms/PGOProfile/prof-verify.ll b/llvm/test/Transforms/PGOProfile/prof-verify.ll
new file mode 100644
index 0000000000000..f4d2b2e51dc59
--- /dev/null
+++ b/llvm/test/Transforms/PGOProfile/prof-verify.ll
@@ -0,0 +1,19 @@
+; Test prof-inject and prof-verify
+
+; RUN: opt -passes=prof-inject %s -S -o - | FileCheck %s --check-prefix=INJECT
+; RUN: not opt -passes=prof-verify %s -S -o - 2>&1 | FileCheck %s --check-prefix=VERIFY
+; RUN: opt -passes=prof-inject,prof-verify %s --disable-output
+
+define void @foo(i32 %i) {
+  %c = icmp eq i32 %i, 0
+  br i1 %c, label %yes, label %no
+yes:
+  ret void
+no:
+  ret void
+}
+
+; INJECT: br i1 %c, label %yes, label %no, !prof !0
+; INJECT: !0 = !{!"branch_weights", i32 429496729, i32 715827882}
+
+; VERIFY: Profile verification failed
\ No newline at end of file

};
} // namespace

bool ProfileInjector::inject() {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can SamplePGO's weight propagation be reused here?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(capturing offline chat) this would be e.g. profi, and yes, we could use it in a number of ways:

  • instead of the perhaps more naive BPI estimator currently used (implicitly) here
  • as the way the BPI estimator handles MD_prof unknown cases (I think these 2 items basically amount to the same thing - improving the BPI estimator)
  • as a numerical validation heuristic (part "2" in the RFC)

@@ -0,0 +1,21 @@
; Test that prof-verify does not modify existing metadata (incl. "unknown")
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Whether overriding unknown MD_prof should probably be guarded by a flag. There are cases the pass can not do a sophisticated weight propagation to infer weight for new branches, but the prof injection pass can.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The expectation is that unknown is inserted by the developer of a pass, so the injector shouldn't override it, since the injector is only here for testing purposes - i.e. its job is to make metadata "be there" for .ll tests that don't originally have it.

@mtrofin mtrofin force-pushed the users/mtrofin/06-30-_pgo_profile_validation branch from 9c22f8c to 99e4436 Compare July 7, 2025 23:17
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
llvm:transforms PGO Profile Guided Optimizations
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants